/*
 * meituan.com Inc.
 * Copyright (c) 2010-2018 All Rights Reserved.
 */
package java.zk.example;

import org.apache.zookeeper.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * <p>
 * 从节点
 * </p>
 * @author LvJing
 * @version $Id:Worker.java v1.0 2018/9/2 上午12:40 LvJing Exp $
 */
public class Worker implements Watcher {

    private static final Logger log = LoggerFactory.getLogger(Worker.class);

    private String serverId = Integer.toHexString(new Random().nextInt());

    private ZooKeeper zk;

    private String hostPort;

    private String status;

    /**
     * 保存正在被该节点执行中的任务，防止任务多次执行
     */
    private final List<String> runningTask;

    /**
     * 从节点执行任务时的线程池
     */
    private ThreadPoolExecutor executor = new ThreadPoolExecutor(10, 5, 0, TimeUnit.MILLISECONDS,
            new ArrayBlockingQueue<>(10));

    public Worker(String hostPort) {
        this.hostPort = hostPort;
        this.runningTask = new ArrayList<>();
    }

    public void startZK() throws IOException {
        this.zk = new ZooKeeper(hostPort, 150000, this);
    }

    public void stopZK() throws InterruptedException {
        this.zk.close();
    }

    @Override
    public void process(WatchedEvent event) {
        log.info(event.toString() + "," + hostPort);
    }

    public void register() {
        this.zk.create("/workers/worker-" + serverId, "Idle".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL, (rc, path, ctx, name) -> {
                    switch (KeeperException.Code.get(rc)) {
                        case CONNECTIONLOSS:
                            register();// 连接丢失，可以简单地再次重试
                            break;
                        case OK:
                            log.info("{} register success", serverId);
                            break;
                        case NODEEXISTS:
                            log.warn("node exist, {}", serverId);
                            break;
                        default:
                            log.error("create worker-{} error", serverId,
                                    KeeperException.create(KeeperException.Code.get(rc)));
                    }
                }, null);
    }

    public void setStatus(String status) {
        this.status = status;// 把传递的status保存至对象本地，防止连接丢失后找不到重试对比的status
        updateStatus(status);
    }

    /**
     * 每个从节点可以分配多个任务，分配给从节点的任务添加在/assign/worker-id/下。
     * 内置一个List存储从节点已经知道该任务分配给它的记录集合，防止任务被从节点多次执行。
     */
    public void getTasks() {
        zk.getChildren("/assign/worker-" + serverId, event -> {
            if (event.getType().equals(Event.EventType.NodeChildrenChanged)) {
                if (!event.getPath().equals("/assign/worker-" + serverId)) {
                    log.warn("任务分配路径改变，path:{}", event.getPath());
                    return;
                }
                log.info("{} 检测到有新任务分配", serverId);
                getTasks();
            }
        }, (rc, path, ctx, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    log.warn("{} 连接丢失，重新查询该节点分配的任务", serverId);
                    getTasks();
                    break;
                case OK:
                    if (children == null || children.size() <= 0)
                        break;
                    executeTask(children);
                default:
                    log.error("获取从节点{}分配任务列表失败", serverId, KeeperException.create(KeeperException.Code.get(rc)));
            }
        }, null);
    }

    private void executeTask(List<String> children) {
        executor.execute(() -> {
            log.info("{} 开始分配任务去执行，任务数量:{}", serverId, children.size());
            synchronized (runningTask) {
                for (String task : children) {
                    if (runningTask.contains(task))
                        continue;
                    log.info("{} 开始执行任务:{}", serverId, task);
                    zk.getData("/assign/worker-" + serverId + "/" + task, false, (rc1, path1, ctx1, data, stat) -> {
                        switch (KeeperException.Code.get(rc1)) {
                            case CONNECTIONLOSS:
                                log.warn("{} 连接丢失，尝试重新获取任务数据", serverId);
                                break;
                            case OK:
                                log.info("任务:{}, 数据:{}", task, new String(data));
                                break;
                            default:
                                log.error("获取任务{}出现异常", task, KeeperException.create(KeeperException.Code.get(rc1)));
                        }
                    }, task);
                }
            }
        });
    }

    private synchronized void updateStatus(String status) {
        // 同步调用，防止多次重试时，如果同时回调，造成多次修改的情况
        if (this.status.equals(status)) {
            this.zk.setData("/workers/worker-" + serverId, status.getBytes(), -1, (rc, path, ctx, stat) -> {
                switch (KeeperException.Code.get(rc)) {
                    case CONNECTIONLOSS:
                        updateStatus((String) ctx);// 连接丢失，再次重试。这里ctx的来源就是setData时最后一个参数
                        break;
                    case OK:
                        log.info("setStatus {} success", status);
                        break;
                    default:
                        log.error("setStatus error", KeeperException.create(KeeperException.Code.get(rc)));
                }
            }, status);
        }
    }

    public static void main(String[] args) throws Exception {
        Worker worker = new Worker("127.0.0.1:2181");
        worker.startZK();
        worker.register();
        worker.setStatus("success");
        worker.getTasks();
        Thread.sleep(60000);
    }
}
